"""
空气质量数据存储模块

功能:
1. 定义空气质量数据模型
2. 提供数据库初始化功能
3. 实现数据批量存储
4. 支持MySQL数据库

使用说明:
1. 需要配置数据库连接URL
2. 支持直接使用pymysql或SQLAlchemy
3. 自动创建数据表结构
"""
from sqlalchemy import create_engine, Column, DateTime, Float, Integer
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import pymysql
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config.db_config import get_db_url

Base = declarative_base()

class AirQualityRecord(Base):
    """
    空气质量数据记录模型
    
    对应数据库表结构:
    - id: 主键
    - timestamp: 时间戳(索引)
    - PM2_5/PM10/SO2/NO2/CO/O3: 污染物浓度
    - AQI: 空气质量指数
    
    注意:
    1. 表名固定为'air_quality'
    2. 使用SQLAlchemy ORM映射
    """
    __tablename__ = 'air_quality'
    id = Column(Integer, primary_key=True)
    timestamp = Column(DateTime, index=True)
    PM2_5 = Column(Float)
    PM10 = Column(Float)
    SO2 = Column(Float)
    NO2 = Column(Float)
    CO = Column(Float)
    O3 = Column(Float)
    AQI = Column(Integer)

def init_db(db_path=None):
    """
    初始化数据库连接
    使用pymysql作为SQLAlchemy的后端驱动
    """
    try:
        if db_path is None:
            db_path = get_db_url()
        engine = create_engine(db_path)
        Base.metadata.create_all(engine)
        return sessionmaker(bind=engine)
    except Exception as e:
        print(f"数据库初始化失败: {str(e)}")
        raise

def save_data_to_mysql(data_records, session_maker=None):
    """
    批量保存空气质量数据到MySQL
    
    参数:
        data_records (list): 数据记录列表，每个元素为字典格式
        session_maker (optional): SQLAlchemy session工厂
        
    返回:
        bool: 保存是否成功
        
    异常处理:
        自动回滚事务并记录错误
    """
    if session_maker is None:
        session_maker = init_db()
    
    session = session_maker()
    try:
        session.bulk_insert_mappings(AirQualityRecord, data_records)
        session.commit()
        return True
    except Exception as e:
        session.rollback()
        print(f"数据保存失败: {str(e)}")
        return False
    finally:
        session.close()
